# -*- coding: utf-8 -*-
# @Time    : 2021/10/20 16:32
# @Author  : huangwei
# @File    : get_cell4.py
# @Software: PyCharm
import cv2
import numpy as np
from skimage import measure

from signature import word_and_box, get_rotate_crop_image
from table_method import draw_lines, minAreaRectBox
from config import args
from get_cell2 import get_lines
from word import TextDetector, TextRecognizer
from method import draw_box

file_path = "label123.png"
img = cv2.imread(file_path)

# 找出所有文本框
text_detector = TextDetector(args)
dt_boxes = text_detector(img)
print("文本框的数量为：", len(dt_boxes))

word_box_img = np.copy(img)
for box in dt_boxes:
    x1, y1 = box[0]
    x2, y2 = box[1]
    x3, y3 = box[2]
    x4, y4 = box[3]

    new_box = [x1, y1, x2, y2, x3, y3, x4, y4]
    draw_box(word_box_img, new_box)
cv2.imwrite("temp_path/word_box.png", word_box_img)

# 找到所有的横线和竖线
input_size = (1024, 1024)
rows_lines, cols_lines = get_lines(img, input_size)


# 找到所有的表格框
def get_table_boxes( img, lines ):
    tmp = np.zeros((img.shape[0], img.shape[1]), dtype='uint8')
    tmp = draw_lines(tmp, lines, color=255, line_width=1)
    tmp_filepath = 'temp_path/tmp.png'
    cv2.imwrite(tmp_filepath, tmp)
    table_labels = measure.label(tmp == 0, connectivity=2)
    regions = measure.regionprops(table_labels)
    rboxes = []
    for region in regions:
        # print("area:", region.bbox_area)
        if region.bbox_area < img.shape[0] * img.shape[1] / 2:
            rbox = minAreaRectBox(region.coords)
            rboxes.append(rbox)

    return rboxes


table_boxes = get_table_boxes(img, rows_lines + cols_lines)
print("表格框数量为：", len(table_boxes))
table_box_img = np.copy(img)
for box in table_boxes:
    draw_box(table_box_img, box)
cv2.imwrite("temp_path/table_box.png", table_box_img)


# 得到分层分列图，移动线条
def split_floors( row_lines, col_lines ):
    # 判断竖线的条数，即分成多少列
    # 从左到右排序了，计算x轴距离
    x_list = []
    for line in col_lines:
        x1, y1, x2, y2 = line
        cx = (x1 + x2) / 2
        x_list.append(cx)

    # 向左合并，即将另一条线左移
    col_num = len(col_lines)

    col_index = [0]

    lx1, _, lx2, _ = col_lines[0]
    lcx = (lx1 + lx2) / 2

    x_index = [lcx]

    for i in range(col_num - 1):
        line1 = col_lines[i]
        line2 = col_lines[i + 1]

        x1, y1, x2, y2 = line1
        cx1 = (x1 + x2) / 2

        x3, y3, x4, y4 = line2
        cx2 = (x3 + x4) / 2

        dis = abs(cx1 - cx2)

        if dis < 5:
            col_index.append(col_index[-1])
            print("左移第{}条线！".format(i + 1))
            x3 = x3 - dis
            x4 = x4 - dis
            line2 = [x3, y3, x4, y4]
            col_lines[i + 1] = line2
        else:
            col_index.append(col_index[-1] + 1)
            x_index.append(cx2)

    # 同理得到y的每一层
    y_index = []
    for line in row_lines:
        x1, y1, x2, y2 = line
        cy = (y1 + y2) / 2
        y_index.append(cy)

    return x_index, y_index, row_lines, col_lines


x_index, y_index, rows_lines, cols_lines = split_floors(rows_lines, cols_lines)

text_recognizer = TextRecognizer(args)


# 得到每一个表格框的信息 [x, y, w, h, txt]
def get_table_box_data( table_boxes, word_boxes, x_index, y_index ):
    box_infos = []
    for table_box in table_boxes:
        x1, y1, x2, y2, x3, y3, x4, y4 = table_box[:8]
        cx = (x1 + x2 + x3 + x4) / 4
        cy = (y1 + y2 + y3 + y4) / 4

        # 左边竖线中心
        left_x = (x1 + x4) / 2
        min_dis = 500
        box_left_index = 0
        # 找到最铁近的竖线
        for m in range(len(x_index)):
            dis = abs(x_index[m] - left_x)
            if dis < min_dis:
                min_dis = dis
                box_left_index = m

        # 同找另外几条线
        # 上边横线
        up_y = (y1 + y2) / 2
        min_dis = 500
        box_up_index = 0
        for m in range(len(y_index)):
            dis = abs(y_index[m] - up_y)
            if dis < min_dis:
                min_dis = dis
                box_up_index = m

        # 右边竖线
        right_x = (x2 + x3) / 2
        min_dis = 500
        box_right_index = 0
        # 找打最铁近的竖线
        for m in range(len(x_index)):
            dis = abs(x_index[m] - right_x)
            if dis < min_dis:
                min_dis = dis
                box_right_index = m

        # 下边横线
        down_y = (y3 + y4) / 2
        min_dis = 500
        box_down_index = 0
        for m in range(len(y_index)):
            dis = abs(y_index[m] - down_y)
            if dis < min_dis:
                min_dis = dis
                box_down_index = m

        result = ''

        temp_list = word_and_box(table_box, word_boxes)
        if len(temp_list) != 0:
            crop_list = []
            for m in range(len(temp_list)):
                img_crop = get_rotate_crop_image(img, temp_list[m])
                crop_list.append(img_crop)
            rec_res = text_recognizer(crop_list)

            # result = ''
            for dd in rec_res:
                result = result + " " + dd[0]

        # for word_box in word_boxes:
        #     wx1, wy1 = word_box[0]
        #     wx2, wy2 = word_box[1]
        #     wx3, wy3 = word_box[2]
        #     wx4, wy4 = word_box[3]
        #
        #     wcx = (wx1 + wx2 + wx3 + wx4) / 4
        #     wcy = (wy1 + wy2 + wy3 + wy4) / 4
        #
        #     if left_x < wcx < right_x and up_y < wcy < down_y:
        #         # 有数据
        #         word_data += 'a'

        location = [box_left_index, box_up_index, box_right_index - box_left_index, box_down_index - box_up_index,
                    result]
        print("location:", location)
        box_infos.append(location)

    return box_infos


box_infos = get_table_box_data(table_boxes, dt_boxes, x_index, y_index)

from openpyxl import Workbook
from openpyxl.styles import Alignment, Side, PatternFill
from openpyxl.styles import Border

sheet_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'O', 'P', 'Q']

book = Workbook()
sheet = book.active

for p in range(len(y_index) - 1):
    for k in range(len(x_index) - 1):
        sheet.cell(p + 1, k + 1).border = Border(top=Side(border_style='thin', color='FF000000'),
                                                 right=Side(border_style='thin', color='FF000000'),
                                                 bottom=Side(border_style='thin', color='FF000000'),
                                                 left=Side(border_style='thin', color='FF000000'))
        sheet.cell(p + 1, k + 1).fill = PatternFill("solid", fgColor="E1FFFF")
        # sheet.cell(p + 1, k + 1).value = 'a'

for tt in box_infos:
    a, b, c, d, data = tt
    A = sheet_list[a]
    B = sheet_list[a + c - 1]
    A1 = '{}{}'.format(A, b + 1)
    B2 = '{}{}'.format(B, b + d)
    sheet.merge_cells('{}:{}'.format(A1, B2))
    sheet['{}'.format(A1)] = data

book.save('temp_path/merging.xlsx')
